#include "Fiber.hpp"
#include "Scheduler.hpp"
#include "Thread.hpp"
#include "../base/Defind.hpp"
#include "../base/BaseLog.hpp"

#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <memory>
#include <sys/ucontext.h>
#include <ucontext.h>

using namespace DBServer::coroutine;

/// @brief 下一个协程的id
static std::atomic<std::size_t> fiber_id = {0};

/// @brief 协程的总数量
static std::atomic<std::size_t> fiber_count = {0};

/// @brief 当前线程主协程
static thread_local Fiber::ptr main_fiber = nullptr;

/// @brief 当前协程的执行协程
static thread_local Fiber* cur_fiber = nullptr;

    /// @brief 获取一个协程的id

static auto get_fiber_id() -> std::size_t { return ++fiber_id; }

/// @brief 协程栈分配器
struct StackAllocer {
    /// @brief 获取一个协程栈
    static auto Get() -> void* {
        void *stack;
        if ((stack = malloc(FIBER_STACK_SIZE)) == nullptr)
        {
            return nullptr;
        }
        return stack;
    }

    /// @brief 将协程栈归还
    static void Back(void *stack) {
        free(stack);
    }
};

/**
 * 
 * 封装ucontext的api，添加log功能并简化的参数的传递
 */
#define swap_ctx(s_fiber, e_fiber)  \
if (swapcontext(&(s_fiber->ctx), &(e_fiber->ctx)) == -1)  \
    error("swap_ctx fail!")

#define get_ctx(fiber)  \
if (getcontext(&(fiber->ctx)) == -1)  \
    error("get_ctx fail!")

#define set_ctx(fiber)  \
if (swapcontext(&(fiber->ctx)) == -1)  \
    error("set_ctx fail!")

#define make_ctx(fiber, func)  \
makecontext(&(fiber->ctx), func, 0)

Fiber::Fiber()
:   state(FiberState::EXEC),
    id(0) 
{
    SetThis(this);
    get_ctx(this);

    ++fiber_count;
    info("%s:main fiber(%ld)\n", Thread::ThreadName().c_str(), id);
}

Fiber::Fiber(std::function<void()> cb_, bool is_schedule_)
:   id(get_fiber_id()),
    state(FiberState::INIT),
    sk(StackAllocer::Get()),
    cb(cb_),
    is_scheduled(is_schedule_)
{
    get_ctx(this);

    if (cb != nullptr) {
        state = FiberState::READY;
    }

    // 设置协程栈的信息
    ctx.uc_link = nullptr;
    ctx.uc_stack.ss_sp = sk;
    ctx.uc_stack.ss_size = FIBER_STACK_SIZE;

    make_ctx(this, &Fiber::MainFunc);

    ++fiber_count;
    info("create a fiber(%ld)\n", id);
}

void Fiber::resume() {
    auto cur = GetThis();
    info("cur_fiber(%d)\nthis_fiber(%d) state(%d)\n", cur->get_id(), this->id, state);
    assert(state == FiberState::READY ||
           ((this == main_fiber.get() ||
             this == Scheduler::GetMainFiber()) &&
            state == FiberState::EXEC));
    assert(cur->get_state() == FiberState::EXEC);
    state = FiberState::EXEC;
    if (cur.get() == main_fiber.get())
    {
        cur->set_state(FiberState::EXEC);
    }
    else if (cur.get() == Scheduler::GetMainFiber())
    {
        cur->set_state(FiberState::EXEC);
    }
    else
    {
        cur->set_state(FiberState::HOLD);
    }

    auto raw_ptr = cur.get();
    cur.reset();

    SetThis(this);
    swap_ctx(raw_ptr, this);
}

void Fiber::reset(std::function<void()> cb_, bool is_schedule_) {
    assert(state == FiberState::TERM 
            || state == FiberState::INIT);
    cb = cb_;

    // 如果没有回调函数就设置为init状态
    state = cb ? FiberState::READY : FiberState::INIT;

    get_ctx(this);
    is_scheduled = is_schedule_;

    ctx.uc_link = nullptr;
    ctx.uc_stack.ss_sp = sk;
    ctx.uc_stack.ss_size = FIBER_STACK_SIZE;

    make_ctx(this, &Fiber::MainFunc);

    info("reset fiber(%ld)\n", id);
}

Fiber::~Fiber() {
    --fiber_count;
    
    if (sk)     // 用户协程有栈
    {
        assert(state == FiberState::TERM);
        StackAllocer::Back(sk);
    }
    else        // 调度协程或者线程的主协程栈
    {
        assert(!cb);
        assert(state == FiberState::EXEC);

        Fiber *cur  = cur_fiber;
        if (cur == this) 
        {
            SetThis(nullptr);
        }
    }
    info("fiber(%ld) exit!\n"
    "total fiber = (%ld)\n", id, fiber_count.load());
}

auto Fiber::GetThis() -> Fiber::ptr {
    if (cur_fiber) 
    {
        return cur_fiber->shared_from_this();
    }

    warn("create a new fiber\n");
    Fiber::ptr main(new Fiber);
    assert(main.get() == cur_fiber);
    main_fiber = main;
    return cur_fiber->shared_from_this();
}

void Fiber::SetThis(Fiber *f) {
    if (f)
    {
        info("running fiber(%ld)\n", f->get_id());
    }
    cur_fiber = f;
}

void Fiber::Yield() {
    auto cur = GetThis();
    FiberState cur_state = cur->get_state();
    debug("cur_state(%d)\n", cur_state);
    assert(cur_state == FiberState::EXEC || cur_state == FiberState::TERM);

    // 如果协程执行完全部的内容就直接退出
    if (cur->get_state() == FiberState::EXEC)
    {
        cur->set_state(FiberState::HOLD);
    }

    if (cur->scheduled())
    {
        auto sche = Scheduler::GetMainFiber();
        SetThis(sche);
        auto raw_ptr = cur.get();
        cur.reset();
        swap_ctx(raw_ptr, sche);
    }
    else 
    {
        SetThis(main_fiber.get());
        auto raw_ptr = cur.get();
        cur.reset();
        swap_ctx(raw_ptr, main_fiber);
    }
}

void Fiber::MainFunc() {
    auto cur = GetThis();
    assert(cur);

    try
    {
        cur->cb();
        cur->cb = nullptr;
        cur->state = FiberState::TERM;
    }
    catch (const std::exception& e)
    {
        error("fiber(%ld):%s\n", cur->get_id(), e.what());
    }

    auto raw_ptr = cur.get();
    cur.reset();
    Yield();

    error("fiber(%ld) not exit!\n", raw_ptr->get_id());
}

auto Fiber::TotalFibers() -> std::uint64_t {
    return fiber_count;
}